import { Context, Next } from 'koa';

/**
 * 内存存储的请求计数器
 */
interface RequestCounter {
  count: number;
  resetTime: number;
}

/**
 * 速率限制存储
 */
class RateLimitStore {
  private store: Map<string, RequestCounter> = new Map();
  private cleanupInterval: NodeJS.Timeout;

  constructor() {
    // 每分钟清理一次过期的记录
    this.cleanupInterval = setInterval(() => {
      this.cleanup();
    }, 60 * 1000);
  }

  /**
   * 获取请求计数
   */
  get(key: string): RequestCounter | undefined {
    const record = this.store.get(key);
    if (record && Date.now() > record.resetTime) {
      this.store.delete(key);
      return undefined;
    }
    return record;
  }

  /**
   * 设置请求计数
   */
  set(key: string, value: RequestCounter): void {
    this.store.set(key, value);
  }

  /**
   * 删除记录
   */
  delete(key: string): void {
    this.store.delete(key);
  }

  /**
   * 清理过期记录
   */
  private cleanup(): void {
    const now = Date.now();
    for (const [key, record] of this.store.entries()) {
      if (now > record.resetTime) {
        this.store.delete(key);
      }
    }
  }

  /**
   * 销毁存储
   */
  destroy(): void {
    if (this.cleanupInterval) {
      clearInterval(this.cleanupInterval);
    }
    this.store.clear();
  }
}

// 全局速率限制存储实例
const rateLimitStore = new RateLimitStore();

/**
 * 速率限制中间件配置
 */
export interface RateLimitOptions {
  max: number;           // 最大请求数
  windowMs: number;      // 时间窗口（毫秒）
  message?: string;      // 超限时的错误消息
  keyGenerator?: (ctx: Context) => string; // 自定义key生成器
  skipSuccessfulRequests?: boolean; // 是否跳过成功的请求
  skipFailedRequests?: boolean;     // 是否跳过失败的请求
}

/**
 * 创建速率限制中间件
 */
export function rateLimitMiddleware(
  max: number,
  windowMs: number,
  options: Partial<RateLimitOptions> = {}
) {
  const config: RateLimitOptions = {
    max,
    windowMs: windowMs * 1000, // 转换为毫秒
    message: '请求过于频繁，请稍后再试',
    keyGenerator: (ctx: Context) => ctx.ip,
    skipSuccessfulRequests: false,
    skipFailedRequests: false,
    ...options
  };

  return async (ctx: Context, next: Next) => {
    const key = config.keyGenerator!(ctx);
    const now = Date.now();
    const resetTime = now + config.windowMs;

    let record = rateLimitStore.get(key);

    if (!record) {
      // 创建新记录
      record = {
        count: 1,
        resetTime
      };
      rateLimitStore.set(key, record);
    } else {
      // 增加计数
      record.count++;
      rateLimitStore.set(key, record);
    }

    // 检查是否超过限制
    if (record.count > config.max) {
      ctx.status = 429;
      ctx.body = {
        success: false,
        message: config.message,
        retryAfter: Math.ceil((record.resetTime - now) / 1000)
      };
      
      // 设置响应头
      ctx.set('X-RateLimit-Limit', config.max.toString());
      ctx.set('X-RateLimit-Remaining', '0');
      ctx.set('X-RateLimit-Reset', Math.ceil(record.resetTime / 1000).toString());
      ctx.set('Retry-After', Math.ceil((record.resetTime - now) / 1000).toString());
      
      return;
    }

    // 设置响应头
    ctx.set('X-RateLimit-Limit', config.max.toString());
    ctx.set('X-RateLimit-Remaining', (config.max - record.count).toString());
    ctx.set('X-RateLimit-Reset', Math.ceil(record.resetTime / 1000).toString());

    await next();

    // 根据配置决定是否跳过计数
    if (
      (config.skipSuccessfulRequests && ctx.status < 400) ||
      (config.skipFailedRequests && ctx.status >= 400)
    ) {
      // 回滚计数
      const currentRecord = rateLimitStore.get(key);
      if (currentRecord && currentRecord.count > 0) {
        currentRecord.count--;
        rateLimitStore.set(key, currentRecord);
      }
    }
  };
}

/**
 * 基于IP的速率限制
 */
export const ipRateLimit = (max: number, windowMs: number) => {
  return rateLimitMiddleware(max, windowMs, {
    keyGenerator: (ctx: Context) => ctx.ip
  });
};

/**
 * 基于用户的速率限制
 */
export const userRateLimit = (max: number, windowMs: number) => {
  return rateLimitMiddleware(max, windowMs, {
    keyGenerator: (ctx: Context) => {
      const userId = (ctx.state.user as any)?.id;
      return userId ? `user:${userId}` : ctx.ip;
    }
  });
};

/**
 * 基于API端点的速率限制
 */
export const endpointRateLimit = (max: number, windowMs: number) => {
  return rateLimitMiddleware(max, windowMs, {
    keyGenerator: (ctx: Context) => `${ctx.ip}:${ctx.path}`
  });
};

/**
 * 销毁速率限制存储（用于测试清理）
 */
export const destroyRateLimitStore = () => {
  rateLimitStore.destroy();
};

export default rateLimitMiddleware;